import os
import json
import numpy as np
import torch
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.manifold import TSNE
from torch.utils.data import DataLoader, TensorDataset
from cnn_model import get_model
from utils import load_and_preprocess_data, SignalDataset

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号


def load_training_history(model_type):
    """加载训练历史"""
    if model_type == "sim":
        history_path = os.path.join("./diagnosis_models/cnn", "sim", "training_history.json")
    elif model_type == "real":
        history_path = os.path.join("./diagnosis_models/cnn", "real", "training_history.json")
    else:
        history_path = os.path.join("./diagnosis_models/cnn", model_type, "training_history.json")

    try:
        with open(history_path, 'r') as f:
            history = json.load(f)
        return history
    except FileNotFoundError:
        print(f"警告：找不到{model_type}的训练历史文件：{history_path}")
        return None
    except json.JSONDecodeError:
        print(f"警告：{model_type}的训练历史文件格式错误：{history_path}")
        return None


def plot_accuracy_curves():
    """绘制准确率曲线对比图"""
    plt.figure(figsize=(12, 6))

    # 加载所有模型类型的训练历史
    model_types = ["vanilla", "wgan", "cgan", "wgan_gp", "sim", "real"]
    colors = ['red', 'green', 'purple', 'orange', 'brown', 'blue']

    for model_type, color in zip(model_types, colors):
        history = load_training_history(model_type)
        if history is not None:
            plt.plot(history['train_acc'], label=f'{model_type} (训练)', color=color, linestyle='-')

    plt.title('不同模型类型的训练集准确率曲线对比')
    plt.xlabel('训练轮次')
    plt.ylabel('准确率 (%)')
    plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left')
    plt.grid(True)
    plt.tight_layout()
    plt.savefig('./comparison_results/accuracy_curves.png', dpi=300, bbox_inches='tight')
    plt.close()


def plot_accuracy_bars():
    """绘制准确率柱状图对比"""
    model_types = ["vanilla", "wgan", "cgan", "wgan_gp", "sim", "real"]
    test_accuracies = []

    for model_type in model_types:
        if model_type == "sim":
            results_path = os.path.join("./diagnosis_models/cnn", "sim", "evaluation_results.json")
        elif model_type == "real":
            results_path = os.path.join("./diagnosis_models/cnn", "real", "evaluation_results.json")
        else:
            results_path = os.path.join("./diagnosis_models/cnn", model_type, "evaluation_results.json")

        try:
            with open(results_path, 'r') as f:
                results = json.load(f)
                # 确保准确率是百分比格式
                accuracy = results['test_accuracy']
                if accuracy < 1:  # 如果是小数形式，转换为百分比
                    accuracy = accuracy * 100
                test_accuracies.append(accuracy)
        except FileNotFoundError:
            print(f"警告：找不到{model_type}的评估结果文件：{results_path}")
            test_accuracies.append(0)
        except json.JSONDecodeError:
            print(f"警告：{model_type}的评估结果文件格式错误：{results_path}")
            test_accuracies.append(0)

    plt.figure(figsize=(10, 6))
    bars = plt.bar(model_types, test_accuracies)

    # 在柱子上添加具体数值
    for bar in bars:
        height = bar.get_height()
        plt.text(bar.get_x() + bar.get_width() / 2., height,
                 f'{height:.2f}%',
                 ha='center', va='bottom')

    plt.title('不同模型类型的测试集准确率对比')
    plt.xlabel('模型类型')
    plt.ylabel('准确率 (%)')
    plt.ylim(0, 100)
    plt.grid(True, axis='y')
    plt.tight_layout()
    plt.savefig('./comparison_results/accuracy_bars.png', dpi=300)
    plt.close()


def plot_tsne_comparison():
    """绘制t-SNE特征分布对比图"""
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model_types = ["vanilla", "wgan", "cgan", "wgan_gp", "sim"]

    plt.figure(figsize=(15, 10))
    for i, model_type in enumerate(model_types, 1):
        plt.subplot(2, 3, i)

        # 加载模型和数据
        if model_type == "sim":
            model_path = os.path.join("./diagnosis_models/cnn", "sim", "best_model.pth")
            use_sim = True
            use_gan = False
            gan_type = None
        else:
            model_path = os.path.join("./diagnosis_models/cnn", model_type, "best_model.pth")
            use_sim = False
            use_gan = True
            gan_type = model_type

        # 加载数据
        X_train, X_test, y_train, y_test = load_and_preprocess_data(
            use_gan=use_gan,
            gan_type=gan_type,
            use_sim=use_sim
        )

        # 创建数据加载器
        test_dataset = SignalDataset(X_test, y_test)
        test_loader = DataLoader(test_dataset, batch_size=64)

        # 加载模型
        model = get_model(num_classes=len(np.unique(y_test)), device=device)
        try:
            # 尝试加载模型权重
            state_dict = torch.load(model_path, map_location=device)
            # 如果模型结构不匹配，打印警告并跳过这个模型
            if not all(k in state_dict for k in model.state_dict().keys()):
                print(f"警告：{model_type}模型的权重结构与当前模型不匹配，跳过此模型")
                continue
            model.load_state_dict(state_dict)
        except Exception as e:
            print(f"警告：加载{model_type}模型时出错：{str(e)}，跳过此模型")
            continue

        model.eval()

        # 提取特征
        features = []
        labels = []
        with torch.no_grad():
            for inputs, targets in test_loader:
                inputs = inputs.to(device)
                if inputs.dim() == 2:
                    inputs = inputs.unsqueeze(1)
                outputs = model(inputs, return_features=True)
                features.append(outputs.cpu().numpy())
                labels.extend(targets.numpy())

        features = np.concatenate(features, axis=0)
        labels = np.array(labels)

        # 应用t-SNE
        tsne = TSNE(n_components=2, random_state=42)
        features_tsne = tsne.fit_transform(features)

        # 绘制散点图
        scatter = plt.scatter(features_tsne[:, 0], features_tsne[:, 1],
                              c=labels, cmap='tab10', alpha=0.6)
        plt.title(f'{model_type} t-SNE特征分布')
        plt.xlabel('t-SNE维度1')
        plt.ylabel('t-SNE维度2')

    plt.tight_layout()
    plt.savefig('./comparison_results/tsne_comparison.png', dpi=300)
    plt.close()


def calculate_prd(real_signal, generated_signal):
    """计算PRD (Percentage Root-mean-square Deviation)"""
    return np.sqrt(np.sum((real_signal - generated_signal) ** 2) / np.sum(real_signal ** 2)) * 100


def calculate_rmse(real_signal, generated_signal):
    """计算RMSE (Root Mean Square Error)"""
    return np.sqrt(np.mean((real_signal - generated_signal) ** 2))


def calculate_mmd(real_samples, generated_samples, sigma=1.0):
    """计算MMD (Maximum Mean Discrepancy)"""

    def gaussian_kernel(x, y, sigma):
        return np.exp(-np.sum((x - y) ** 2) / (2 * sigma ** 2))

    n_real = len(real_samples)
    n_gen = len(generated_samples)

    # 计算核矩阵
    k_xx = np.zeros((n_real, n_real))
    k_yy = np.zeros((n_gen, n_gen))
    k_xy = np.zeros((n_real, n_gen))

    for i in range(n_real):
        for j in range(n_real):
            k_xx[i, j] = gaussian_kernel(real_samples[i], real_samples[j], sigma)

    for i in range(n_gen):
        for j in range(n_gen):
            k_yy[i, j] = gaussian_kernel(generated_samples[i], generated_samples[j], sigma)

    for i in range(n_real):
        for j in range(n_gen):
            k_xy[i, j] = gaussian_kernel(real_samples[i], generated_samples[j], sigma)

    # 计算MMD
    mmd = np.mean(k_xx) + np.mean(k_yy) - 2 * np.mean(k_xy)
    return mmd


def calculate_metrics():
    """计算并保存评估指标"""
    metrics = {}
    gan_types = ["vanilla", "wgan", "cgan", "wgan_gp"]

    # 加载原始测试数据作为参考
    test_data_path = os.path.join("./cwru_prepro", "test")
    real_samples = []
    for file in os.listdir(test_data_path):
        if file.endswith('.npy'):
            data = np.load(os.path.join(test_data_path, file))
            real_samples.append(data)
    real_samples = np.concatenate(real_samples)

    for gan_type in gan_types:
        # 加载GAN生成的样本（使用最佳模型生成的样本）
        gan_data_path = os.path.join("./cwru_gan", gan_type, "best_model")
        if not os.path.exists(gan_data_path):
            print(f"警告：找不到{gan_type}的最佳模型生成的样本，将使用最终模型生成的样本")
            gan_data_path = os.path.join("./cwru_gan", gan_type)
        
        gen_samples = []
        for file in os.listdir(gan_data_path):
            if file.endswith('.npy'):
                data = np.load(os.path.join(gan_data_path, file))
                gen_samples.append(data)
        gen_samples = np.concatenate(gen_samples)

        # 计算指标
        prd_values = []
        rmse_values = []
        for i in range(min(len(real_samples), len(gen_samples))):
            prd_values.append(calculate_prd(real_samples[i], gen_samples[i]))
            rmse_values.append(calculate_rmse(real_samples[i], gen_samples[i]))

        mmd_value = calculate_mmd(real_samples[:100], gen_samples[:100])  # 使用前100个样本计算MMD

        metrics[gan_type] = {
            'PRD_mean': np.mean(prd_values),
            'PRD_std': np.std(prd_values),
            'RMSE_mean': np.mean(rmse_values),
            'RMSE_std': np.std(rmse_values),
            'MMD': mmd_value
        }

    # 保存指标
    with open('./comparison_results/metrics.json', 'w') as f:
        json.dump(metrics, f, indent=4)

    # 绘制指标对比图
    plt.figure(figsize=(15, 5))

    # PRD对比
    plt.subplot(1, 3, 1)
    prd_means = [metrics[gan_type]['PRD_mean'] for gan_type in gan_types]
    prd_stds = [metrics[gan_type]['PRD_std'] for gan_type in gan_types]
    plt.bar(gan_types, prd_means)
    plt.title('PRD对比')
    plt.xlabel('GAN类型')
    plt.ylabel('PRD值')

    # RMSE对比
    plt.subplot(1, 3, 2)
    rmse_means = [metrics[gan_type]['RMSE_mean'] for gan_type in gan_types]
    rmse_stds = [metrics[gan_type]['RMSE_std'] for gan_type in gan_types]
    plt.bar(gan_types, rmse_means)
    plt.title('RMSE对比')
    plt.xlabel('GAN类型')
    plt.ylabel('RMSE值')

    # MMD对比
    plt.subplot(1, 3, 3)
    mmd_values = [metrics[gan_type]['MMD'] for gan_type in gan_types]
    plt.bar(gan_types, mmd_values)
    plt.title('MMD对比')
    plt.xlabel('GAN类型')
    plt.ylabel('MMD值')

    plt.tight_layout()
    plt.savefig('./comparison_results/metrics_comparison.png', dpi=300)
    plt.close()


def plot_loss_curves():
    """绘制不同GAN模型的损失变化曲线"""
    plt.figure(figsize=(15, 10))
    
    # 加载所有模型类型的训练历史
    model_types = ["vanilla", "wgan", "cgan", "wgan_gp"]
    colors = ['red', 'green', 'purple', 'orange']
    
    for model_type, color in zip(model_types, colors):
        history_path = os.path.join("./models", model_type, "training_history.json")
        try:
            with open(history_path, 'r') as f:
                history = json.load(f)
                
            # 绘制生成器损失
            plt.subplot(2, 2, 1)
            plt.plot(history['g_loss'], label=f'{model_type}', color=color)
            plt.title('生成器损失变化')
            plt.xlabel('训练轮次')
            plt.ylabel('损失值')
            plt.legend()
            plt.grid(True)
            
            # 绘制判别器损失
            plt.subplot(2, 2, 2)
            plt.plot(history['d_loss'], label=f'{model_type}', color=color)
            plt.title('判别器损失变化')
            plt.xlabel('训练轮次')
            plt.ylabel('损失值')
            plt.legend()
            plt.grid(True)
            
        except FileNotFoundError:
            print(f"警告：找不到{model_type}的训练历史文件：{history_path}")
        except json.JSONDecodeError:
            print(f"警告：{model_type}的训练历史文件格式错误：{history_path}")
    
    plt.tight_layout()
    plt.savefig('./comparison_results/loss_curves.png', dpi=300)
    plt.close()


def calculate_signal_quality_metrics(real_signal, generated_signal):
    """计算信号质量指标"""
    # 计算信噪比 (SNR)
    signal_power = np.mean(real_signal ** 2)
    noise_power = np.mean((real_signal - generated_signal) ** 2)
    snr = 10 * np.log10(signal_power / noise_power) if noise_power > 0 else float('inf')
    
    # 计算峰值信噪比 (PSNR)
    max_signal = np.max(np.abs(real_signal))
    mse = np.mean((real_signal - generated_signal) ** 2)
    psnr = 20 * np.log10(max_signal / np.sqrt(mse)) if mse > 0 else float('inf')
    
    # 计算结构相似性指数 (SSIM)
    from skimage.metrics import structural_similarity as ssim
    ssim_value = ssim(real_signal, generated_signal, data_range=max_signal)
    
    return {
        'SNR': snr,
        'PSNR': psnr,
        'SSIM': ssim_value
    }


def plot_signal_quality_metrics():
    """绘制信号质量指标对比图"""
    metrics = {}
    gan_types = ["vanilla", "wgan", "cgan", "wgan_gp"]
    
    # 加载原始测试数据作为参考
    test_data_path = os.path.join("./cwru_prepro", "test")
    real_samples = []
    for file in os.listdir(test_data_path):
        if file.endswith('.npy'):
            data = np.load(os.path.join(test_data_path, file))
            real_samples.append(data)
    real_samples = np.concatenate(real_samples)
    
    for gan_type in gan_types:
        # 加载GAN生成的样本
        gan_data_path = os.path.join("./cwru_gan", gan_type)
        gen_samples = []
        for file in os.listdir(gan_data_path):
            if file.endswith('.npy'):
                data = np.load(os.path.join(gan_data_path, file))
                gen_samples.append(data)
        gen_samples = np.concatenate(gen_samples)
        
        # 计算信号质量指标
        snr_values = []
        psnr_values = []
        ssim_values = []
        
        for i in range(min(len(real_samples), len(gen_samples))):
            quality_metrics = calculate_signal_quality_metrics(real_samples[i], gen_samples[i])
            snr_values.append(quality_metrics['SNR'])
            psnr_values.append(quality_metrics['PSNR'])
            ssim_values.append(quality_metrics['SSIM'])
        
        metrics[gan_type] = {
            'SNR_mean': np.mean(snr_values),
            'SNR_std': np.std(snr_values),
            'PSNR_mean': np.mean(psnr_values),
            'PSNR_std': np.std(psnr_values),
            'SSIM_mean': np.mean(ssim_values),
            'SSIM_std': np.std(ssim_values)
        }
    
    # 绘制信号质量指标对比图
    plt.figure(figsize=(15, 5))
    
    # SNR对比
    plt.subplot(1, 3, 1)
    snr_means = [metrics[gan_type]['SNR_mean'] for gan_type in gan_types]
    snr_stds = [metrics[gan_type]['SNR_std'] for gan_type in gan_types]
    plt.bar(gan_types, snr_means, yerr=snr_stds, capsize=5)
    plt.title('SNR对比')
    plt.xlabel('GAN类型')
    plt.ylabel('SNR (dB)')
    
    # PSNR对比
    plt.subplot(1, 3, 2)
    psnr_means = [metrics[gan_type]['PSNR_mean'] for gan_type in gan_types]
    psnr_stds = [metrics[gan_type]['PSNR_std'] for gan_type in gan_types]
    plt.bar(gan_types, psnr_means, yerr=psnr_stds, capsize=5)
    plt.title('PSNR对比')
    plt.xlabel('GAN类型')
    plt.ylabel('PSNR (dB)')
    
    # SSIM对比
    plt.subplot(1, 3, 3)
    ssim_means = [metrics[gan_type]['SSIM_mean'] for gan_type in gan_types]
    ssim_stds = [metrics[gan_type]['SSIM_std'] for gan_type in gan_types]
    plt.bar(gan_types, ssim_means, yerr=ssim_stds, capsize=5)
    plt.title('SSIM对比')
    plt.xlabel('GAN类型')
    plt.ylabel('SSIM值')
    
    plt.tight_layout()
    plt.savefig('./comparison_results/signal_quality_metrics.png', dpi=300)
    plt.close()
    
    # 保存指标
    with open('./comparison_results/signal_quality_metrics.json', 'w') as f:
        json.dump(metrics, f, indent=4)


def plot_training_time():
    """绘制不同GAN模型的训练时间对比"""
    plt.figure(figsize=(10, 6))
    
    model_types = ["vanilla", "wgan", "cgan", "wgan_gp"]
    training_times = []
    
    for model_type in model_types:
        history_path = os.path.join("./models", model_type, "training_history.json")
        try:
            with open(history_path, 'r') as f:
                history = json.load(f)
                training_times.append(history.get('training_time', 0))
        except (FileNotFoundError, json.JSONDecodeError):
            training_times.append(0)
    
    plt.bar(model_types, training_times)
    plt.title('不同GAN模型的训练时间对比')
    plt.xlabel('GAN类型')
    plt.ylabel('训练时间 (秒)')
    plt.grid(True, axis='y')
    
    # 在柱子上添加具体数值
    for i, v in enumerate(training_times):
        plt.text(i, v, f'{v:.1f}s', ha='center', va='bottom')
    
    plt.tight_layout()
    plt.savefig('./comparison_results/training_time.png', dpi=300)
    plt.close()


def main():
    # 创建结果保存目录
    os.makedirs('./comparison_results', exist_ok=True)

    # 绘制准确率曲线
    print("绘制准确率曲线...")
    plot_accuracy_curves()

    # 绘制准确率柱状图
    print("绘制准确率柱状图...")
    plot_accuracy_bars()

    # 绘制t-SNE对比图
    print("绘制t-SNE特征分布对比图...")
    plot_tsne_comparison()

    # 绘制损失变化曲线
    print("绘制损失变化曲线...")
    plot_loss_curves()

    # 计算评估指标
    print("计算评估指标...")
    calculate_metrics()

    # 计算信号质量指标
    print("计算信号质量指标...")
    plot_signal_quality_metrics()

    # 绘制训练时间对比
    print("绘制训练时间对比...")
    plot_training_time()

    print("对比分析完成！结果保存在 ./comparison_results 目录下")


if __name__ == "__main__":
    main()